package com.xiashitech.agent.instrumentation.db.resultset;

import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.instrumentation.base.AbstractAttributesExtractor;
import com.xiashitech.agent.instrumentation.db.common.XSResultSizeContext;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.instrumentation.utils.MachineAttributesUtil;
import com.xiashitech.agent.utils.SessionUtil;
import com.xiashitech.interfaces.agent.dto.SystemInfoDTO;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.Baggage;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.common.AttributesBuilder;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import org.jetbrains.annotations.Nullable;

public class XSResultSetAttributesExtractor<Request, Void> extends AbstractAttributesExtractor<Request, Void> {
    public String getAttributesExtractorName() {
        return XSResultSetAttributesExtractor.class.getName();
    }

    public static <Request, Void> XSResultSetAttributesExtractor<Request, Void> create() {
        return new XSResultSetAttributesExtractor();
    }

    @Override
    public void onXSStart(AttributesBuilder attributesBuilder, Context context, Request request) {
    }

    @Override
    public void onXSEnd(AttributesBuilder attributesBuilder, Context context, Request request, @Nullable Void _void, @Nullable Throwable throwable) {
        Baggage baggage = SessionUtil.syncBaggage();
        baggage = BaggageUtil.mergeBaggageFromContext(baggage,context).build();
        XSResultSizeContext.Request requestContext = (XSResultSizeContext.Request)request;
        attributesBuilder.put("xs.db.records.size",requestContext.getSize());
        attributesBuilder.put("xs.db.sql",requestContext.getSql());
        attributesBuilder.put("xs.systemId", StartConfig.getParamValueFromKey(StartConfig.systemId));
        if(SystemInfoDTO.getConfig() != null) {
            attributesBuilder.put("xs.systemReg", SystemInfoDTO.getConfig().getSystemExist());
        }
        attributesBuilder.put("xs.sourceSystemId", baggage.getEntryValue("xsAuthSystemId") == null ? StartConfig.getParamValueFromKey(StartConfig.systemId):baggage.getEntryValue("xsAuthSystemId"));
        MachineAttributesUtil.fillMachineAttributes(attributesBuilder);
    }
}
